/*
 * Wazuh Vulnerability scanner - Scan Orchestrator
 * Copyright (C) 2015, Wazuh Inc.
 * March 11, 2024.
 *
 * This program is free software; you can redistribute it
 * and/or modify it under the terms of the GNU General Public
 * License (version 2) as published by the FSF - Free Software
 * Foundation.
 */

#ifndef _CLEAN_AGENT_INVENTORY_HPP
#define _CLEAN_AGENT_INVENTORY_HPP

#include "chainOfResponsability.hpp"
#include "indexerConnector.hpp"
#include "inventorySync.hpp"
#include "loggerHelper.h"
#include "scanContext.hpp"
#include "socketDBWrapper.hpp"
#include "vulnerabilityScanner.hpp"
#include "wazuhDBQueryBuilder.hpp"

/**
 * @brief A class for handling data cleanup operations in the inventory database and publishing the changes to the
 * indexer.
 * This class is responsible for deleting all the inventory entries for a given agent and affected component type.
 * It is also responsible to call a sub-orchestration to publish the changes to the indexer.
 * It receives the scan context and the inventory database and returns the scan context with the inventory cleaned.
 * The affected component type can be Agent, Os or Package.
 * Agent: Delete all inventories for the agent.
 * Os: Delete inventory of the OS for the agent.
 * Package: Delete inventory of the package for the agent.
 *
 * @tparam TIndexerConnector Indexer connector.
 * @tparam TScanContext Scan context.
 */
template<typename TScanContext = ScanContext,
         typename TAbstractHandler = AbstractHandler<std::shared_ptr<TScanContext>>>
class TCleanAgentInventory final
    : public AbstractHandler<std::shared_ptr<TScanContext>>
    , public TInventorySync<TScanContext>
{
private:
    std::shared_ptr<TAbstractHandler> m_subOrchestration;

public:
    /**
     * @brief Class constructor.
     * @param inventoryDatabase Inventory database.
     * @param subOrchestration Sub-orchestration to publish changes to the indexer.
     */
    explicit TCleanAgentInventory(Utils::RocksDBWrapper& inventoryDatabase,
                                  std::shared_ptr<TAbstractHandler> subOrchestration)
        : m_subOrchestration(std::move(subOrchestration))
        , TInventorySync<TScanContext>(inventoryDatabase)
    {
    }

    /**
     * @brief Handles request and passes control to the next step of the chain.
     *
     * @param data Scan context.
     * @return std::shared_ptr<TScanContext> Abstract handler.
     */
    std::shared_ptr<TScanContext> handleRequest(std::shared_ptr<TScanContext> data) override
    {
        std::string agentKey;
        agentKey.append(data->agentNodeName());
        agentKey.append("_");
        agentKey.append(data->agentId());
        agentKey.append("_");

        data->m_isInventoryEmpty = true;

        auto deleteAll = [this](std::shared_ptr<TScanContext> data,
                                const std::string& key,
                                const std::vector<AffectedComponentType>& types)
        {
            for (const auto& type : types)
            {
                for (const auto& dbQuery :
                     TInventorySync<TScanContext>::m_inventoryDatabase.seek(key, AFFECTED_COMPONENT_COLUMNS.at(type)))
                {
                    auto listCve = Utils::split(dbQuery.second.ToString(), ',');
                    auto context = std::make_shared<TScanContext>();
                    for (const auto& cve : listCve)
                    {
                        data->m_isInventoryEmpty = false;
                        std::string elementKey;
                        elementKey.append(dbQuery.first);
                        elementKey.append("_");
                        elementKey.append(cve);

                        context->m_elements.emplace(cve,
                                                    TInventorySync<TScanContext>::buildElement("DELETED", elementKey));
                    }
                    m_subOrchestration->handleRequest(std::move(context));
                    logDebug2(WM_VULNSCAN_LOGTAG,
                              "Deleting %s agent vulnerabilities key: %s",
                              AFFECTED_COMPONENT_COLUMNS.at(type).c_str(),
                              dbQuery.first.c_str());
                    TInventorySync<TScanContext>::m_inventoryDatabase.delete_(dbQuery.first,
                                                                              AFFECTED_COMPONENT_COLUMNS.at(type));
                }
            }
        };

        // If the affected type is Agent, delete all inventories for the agent.
        if (data->affectedComponentType() == AffectedComponentType::Agent)
        {
            deleteAll(data, agentKey, {AffectedComponentType::Os, AffectedComponentType::Package});
        }
        // Delete all entries for the affected type, used for the integrity clear.
        else
        {
            deleteAll(data, agentKey, {data->affectedComponentType()});
        }

        // If the affected type is Os, delete the initial scan data.
        if (data->affectedComponentType() == AffectedComponentType::Os ||
            data->affectedComponentType() == AffectedComponentType::Agent)
        {
            TInventorySync<TScanContext>::m_inventoryDatabase.delete_(data->agentId().data(), OS_INITIAL_SCAN);
        }

        return AbstractHandler<std::shared_ptr<TScanContext>>::handleRequest(std::move(data));
    }
};

using CleanAgentInventory = TCleanAgentInventory<>;

#endif // _CLEAN_AGENT_INVENTORY_HPP
